//
//  AESUtils.m
//  AESUtilsDemo
//
//  Created by ilongge on 2021/7/29.
//

#import "AESUtils.h"
#import <CommonCrypto/CommonCrypto.h>
#import <CommonCrypto/CommonDigest.h>
@implementation AESUtils

+ (NSData *)encryptWithKey:(NSString *)encryptKey andIvKey:(NSString *)ivKey andData:(NSData *)encryptData
{
    return [AESUtils encryptWithKey:encryptKey andIvKey:ivKey andMode:AES_Mode_ECB andData:encryptData];
}

+ (NSString *)encryptWithKey:(NSString *)encryptKey andIvKey:(NSString *)ivKey andText:(NSString *)encryptText
{
    return [AESUtils encryptWithKey:encryptKey andIvKey:ivKey andMode:AES_Mode_ECB andText:encryptText];
}

+ (NSData *)decryptWithKey:(NSString *)decryptKey andIvKey:(NSString *)ivKey andData:(NSData *)decryptData
{
    return [AESUtils decryptWithKey:decryptKey andIvKey:ivKey andMode:AES_Mode_ECB andData:decryptData];
}

+ (NSString *)decryptWithKey:(NSString *)decryptKey andIvKey:(NSString *)ivKey andText:(NSString *)decryptText
{
    return [AESUtils decryptWithKey:decryptKey andIvKey:ivKey andMode:AES_Mode_ECB andText:decryptText];
}

+ (NSData *)encryptWithKey:(NSString *)encryptKey andIvKey:(NSString *)ivKey andMode:(AES_Mode)mode andData:(NSData *)encryptData
{
    encryptKey = [AESUtils checkAESKey:encryptKey];
    //16位偏移，仅限于CBC模式需要
    NSData *initVector;
    if (mode == AES_Mode_CBC) {
        ivKey = [AESUtils checkAESKey:ivKey];
        initVector = [ivKey dataUsingEncoding:NSUTF8StringEncoding];
    }
    //公钥
    char keyPtr[kCCKeySizeAES128+1];
    memset(keyPtr, 0, sizeof(keyPtr));
    [encryptKey getCString:keyPtr maxLength:sizeof(keyPtr) encoding:NSUTF8StringEncoding];
    //数据长度
    NSUInteger dataLength = encryptData.length;
    //加密输出缓冲区大小
    size_t bufferSize = dataLength + kCCBlockSizeAES128;
    //加密输出缓冲区
    void *buffer = malloc(bufferSize);
    //实际输出大小
    size_t numBytesEncrypted = 0;
    CCCryptorStatus cryptStatus = CCCrypt(kCCEncrypt,//kCCEncrypt代表加密
                                          kCCAlgorithmAES,//加密算法
                                          mode == AES_Mode_CBC?kCCOptionPKCS7Padding:kCCOptionPKCS7Padding|kCCOptionECBMode,//CBC -> PKCS7Padding，ECB -> kCCOptionPKCS7Padding|kCCOptionECBMode
                                          keyPtr,
                                          kCCBlockSizeAES128,//密钥长度128
                                          mode == AES_Mode_CBC?initVector.bytes:NULL,//偏移字符串,ECB模式传NULL
                                          encryptData.bytes,//编码内容
                                          dataLength,
                                          buffer,
                                          bufferSize,
                                          &numBytesEncrypted);
    if (cryptStatus == kCCSuccess) {
        return [NSData dataWithBytesNoCopy:buffer length:numBytesEncrypted];
    }
    free(buffer);
    return nil;
}

+ (NSString *)encryptWithKey:(NSString *)encryptKey andIvKey:(NSString *)ivKey andMode:(AES_Mode)mode andText:(NSString *)encryptText
{
    encryptKey = [AESUtils checkAESKey:encryptKey];
    ivKey = [AESUtils checkAESKey:ivKey];
    NSData *data = [encryptText dataUsingEncoding:NSUTF8StringEncoding];
    //对数据进行加密
    NSData *result = [AESUtils encryptWithKey:encryptKey andIvKey:ivKey andMode:mode andData:data];
    //转换为2进制字符串
    if (result && result.length > 0) {
        Byte *datas = (Byte*)[result bytes];
        NSMutableString *output = [NSMutableString stringWithCapacity:result.length * 2];
        for(int i = 0; i < result.length; i++){
            [output appendFormat:@"%02x", datas[i]];
        }
        return output;
    }
    return nil;
}

+ (NSData *)decryptWithKey:(NSString *)decryptKey andIvKey:(NSString *)ivKey andMode:(AES_Mode)mode andData:(NSData *)decryptData
{
    decryptKey = [AESUtils checkAESKey:decryptKey];
    //16位偏移，仅限于CBC模式需要
    NSData *initVector;
    if (mode == AES_Mode_CBC) {
        ivKey = [AESUtils checkAESKey:ivKey];
        initVector = [ivKey dataUsingEncoding:NSUTF8StringEncoding];
    }
    //公钥
    char keyPtr[kCCKeySizeAES128+1];
    memset(keyPtr, 0, sizeof(keyPtr));
    [decryptKey getCString:keyPtr maxLength:sizeof(keyPtr) encoding:NSUTF8StringEncoding];
    //数据长度
    NSUInteger dataLength = decryptData.length;
    //加密输出缓冲区大小
    size_t bufferSize = dataLength + kCCBlockSizeAES128;
    //加密输出缓冲区
    void *buffer = malloc(bufferSize);
    //实际输出大小
    size_t numBytesDecrypted = 0;
    CCCryptorStatus cryptStatus = CCCrypt(kCCDecrypt,//kCCDecrypt代表解密
                                          kCCAlgorithmAES,//加密算法
                                          mode == AES_Mode_CBC?kCCOptionPKCS7Padding:kCCOptionPKCS7Padding|kCCOptionECBMode,//CBC -> PKCS7Padding，ECB -> kCCOptionPKCS7Padding|kCCOptionECBMode
                                          keyPtr,
                                          kCCBlockSizeAES128,//密钥长度128
                                          mode == AES_Mode_CBC?initVector.bytes:NULL,//偏移字符串,ECB模式传NULL
                                          decryptData.bytes,//编码内容
                                          dataLength,
                                          buffer,
                                          bufferSize,
                                          &numBytesDecrypted);
    if (cryptStatus == kCCSuccess) {
        return [NSData dataWithBytesNoCopy:buffer length:numBytesDecrypted];
    }
    free(buffer);
    return nil;
}

+ (NSString *)decryptWithKey:(NSString *)decryptKey andIvKey:(NSString *)ivKey andMode:(AES_Mode)mode andText:(NSString *)decryptText
{
    decryptKey = [AESUtils checkAESKey:decryptKey];
    ivKey = [AESUtils checkAESKey:ivKey];
    //转换为2进制Data
    NSMutableData *data = [NSMutableData dataWithCapacity:decryptText.length / 2];
    unsigned char whole_byte;
    char byte_chars[3] = {'\0','\0','\0'};
    int i;
    for (i=0; i < [decryptText length] / 2; i++) {
        byte_chars[0] = [decryptText characterAtIndex:i*2];
        byte_chars[1] = [decryptText characterAtIndex:i*2+1];
        whole_byte = strtol(byte_chars, NULL, 16);
        [data appendBytes:&whole_byte length:1];
    }
    
    //对数据进行解密
    NSData* result = [AESUtils decryptWithKey:decryptKey andIvKey:ivKey andMode:mode andData:data];
    if (result && result.length > 0) {
        return [[NSString alloc] initWithData:result encoding:NSUTF8StringEncoding];
    }
    return nil;
}

+ (NSString *)checkAESKey:(NSString *)AES_Key
{
    if (AES_Key.length < 16) {
        NSMutableArray *array = [NSMutableArray arrayWithArray:@[AES_Key]];
        while (array.count <= 16 - AES_Key.length) {
            [array insertObject:@"0" atIndex:0];
        }
        AES_Key = [array componentsJoinedByString:@""];
    }
    return AES_Key;
}
@end

